/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.gaia;



import com.bff.gaia.unified.runners.core.construction.PipelineResources;

import com.bff.gaia.unified.sdk.Pipeline;

import com.bff.gaia.unified.vendor.guava.com.google.common.annotations.VisibleForTesting;

import com.bff.gaia.unified.vendor.guava.com.google.common.base.MoreObjects;

import com.bff.gaia.api.common.JobExecutionResult;

import com.bff.gaia.api.java.ExecutionEnvironment;

import com.bff.gaia.unified.sdk.values.PCollection;

import com.bff.gaia.runtime.jobgraph.JobGraph;

import com.bff.gaia.streaming.api.environment.StreamExecutionEnvironment;

import org.slf4j.Logger;

import org.slf4j.LoggerFactory;



import static com.bff.gaia.unified.vendor.guava.com.google.common.base.Preconditions.checkNotNull;



/**

 * The class that instantiates and manages the execution of a given job. Depending on if the job is

 * a Streaming or Batch processing one, it creates the adequate execution environment ({@link

 * ExecutionEnvironment} or {@link StreamExecutionEnvironment}), the necessary {@link

 * GaiaPipelineTranslator} ({@link GaiaBatchPipelineTranslator} or {@link

 * GaiaStreamingPipelineTranslator}) to transform the Unified job into a Gaia one, and executes the

 * (translated) job.

 */

class GaiaPipelineExecutionEnvironment {



  private static final Logger LOG =

      LoggerFactory.getLogger(GaiaPipelineExecutionEnvironment.class);



  private final GaiaPipelineOptions options;



  /**

   * The Gaia Batch execution environment. This is instantiated to either a {@link

   * com.bff.gaia.api.java.CollectionEnvironment}, a {@link

   * com.bff.gaia.api.java.LocalEnvironment} or a {@link

   * com.bff.gaia.api.java.RemoteEnvironment}, depending on the configuration options.

   */

  private ExecutionEnvironment gaiaBatchEnv;



  /**

   * The Gaia Streaming execution environment. This is instantiated to either a {@link

   * com.bff.gaia.streaming.api.environment.LocalStreamEnvironment} or a {@link

   * com.bff.gaia.streaming.api.environment.RemoteStreamEnvironment}, depending on the

   * configuration options, and more specifically, the url of the master.

   */

  private StreamExecutionEnvironment gaiaStreamEnv;



  /**

   * Creates a {@link GaiaPipelineExecutionEnvironment} with the user-specified parameters in the

   * provided {@link GaiaPipelineOptions}.

   *

   * @param options the user-defined pipeline options.

   */

  GaiaPipelineExecutionEnvironment(GaiaPipelineOptions options) {

    this.options = checkNotNull(options);

  }



  /**

   * Depending on if the job is a Streaming or a Batch one, this method creates the necessary

   * execution environment and pipeline translator, and translates the {@link

   * PCollection} program into a {@link

   * com.bff.gaia.api.java.DataSet} or {@link

   * com.bff.gaia.streaming.api.datastream.DataStream} one.

   */

  public void translate(Pipeline pipeline) {

    this.gaiaBatchEnv = null;

    this.gaiaStreamEnv = null;



    final boolean hasUnboundedOutput =

        PipelineTranslationModeOptimizer.hasUnboundedOutput(pipeline);

    if (hasUnboundedOutput) {

      LOG.info("Found unbounded PCollection. Switching to streaming execution.");

      options.setStreaming(true);

    }



    // Staged files need to be set before initializing the execution environments

    prepareFilesToStageForRemoteClusterExecution(options);



    GaiaPipelineTranslator translator;

    if (options.isStreaming()) {

      this.gaiaStreamEnv =

          GaiaExecutionEnvironments.createStreamExecutionEnvironment(

              options, options.getFilesToStage());

      if (hasUnboundedOutput && !gaiaStreamEnv.getCheckpointConfig().isCheckpointingEnabled()) {

        LOG.warn(

            "UnboundedSources present which rely on checkpointing, but checkpointing is disabled.");

      }

      translator = new GaiaStreamingPipelineTranslator(gaiaStreamEnv, options);

    } else {

      this.gaiaBatchEnv =

          GaiaExecutionEnvironments.createBatchExecutionEnvironment(

              options, options.getFilesToStage());

      translator = new GaiaBatchPipelineTranslator(gaiaBatchEnv, options);

    }



    // Transform replacements need to receive the finalized PipelineOptions

    // including execution mode (batch/streaming) and parallelism.

    pipeline.replaceAll(GaiaTransformOverrides.getDefaultOverrides(options));



    translator.translate(pipeline);

  }



  /**

   * Local configurations work in the same JVM and have no problems with improperly formatted files

   * on classpath (eg. directories with .class files or empty directories). Prepare files for

   * staging only when using remote cluster (passing the master address explicitly).

   */

  private static void prepareFilesToStageForRemoteClusterExecution(GaiaPipelineOptions options) {

    if (!options.getGaiaMaster().matches("\\[auto\\]|\\[collection\\]|\\[local\\]")) {

      options.setFilesToStage(

          PipelineResources.prepareFilesForStaging(

              options.getFilesToStage(),

              MoreObjects.firstNonNull(

                  options.getTempLocation(), System.getProperty("java.io.tmpdir"))));

    }

  }



  /** Launches the program execution. */

  public JobExecutionResult executePipeline() throws Exception {

    final String jobName = options.getJobName();



    if (gaiaBatchEnv != null) {

      return gaiaBatchEnv.execute(jobName);

    } else if (gaiaStreamEnv != null) {

      return gaiaStreamEnv.execute(jobName);

    } else {

      throw new IllegalStateException("The Pipeline has not yet been translated.");

    }

  }



  /**

   * Retrieves the generated JobGraph which can be submitted against the cluster. For testing

   * purposes.

   */

  @VisibleForTesting

  JobGraph getJobGraph(Pipeline p) {

    translate(p);

    return gaiaStreamEnv.getStreamGraph().getJobGraph();

  }



  @VisibleForTesting

  ExecutionEnvironment getBatchExecutionEnvironment() {

    return gaiaBatchEnv;

  }



  @VisibleForTesting

  StreamExecutionEnvironment getStreamExecutionEnvironment() {

    return gaiaStreamEnv;

  }

}